package com.etocrm.saas.lib.base.resource.config;

import com.etocrm.saas.lib.base.common.result.Result;
import com.etocrm.saas.lib.base.common.result.ResultCodeEnum;
import com.etocrm.saas.lib.base.resource.Util.IpAddressUtil;
import lombok.extern.slf4j.Slf4j;
import net.jodah.expiringmap.ExpirationPolicy;
import net.jodah.expiringmap.ExpiringMap;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import javax.servlet.http.HttpServletRequest;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 采用ExpiringMap处理接口请求次数限制
 */
@Aspect
@Component
@Slf4j
public class LimitRequestAspect {

    private static ConcurrentHashMap<String, ExpiringMap<String, Integer>> book = new ConcurrentHashMap<>();
    private static ConcurrentHashMap<String, ExpiringMap<String, Integer>> bookMap1 = new ConcurrentHashMap<>();
    private static ConcurrentHashMap<String, ExpiringMap<String, Integer>> bookMap2 = new ConcurrentHashMap<>();

    // 定义切点
    // 让所有有@LimitRequest注解的方法都执行切面方法
    @Pointcut("@annotation(limitRequest)")
    public void excudeService(LimitRequest limitRequest) {
    }

    @Around("excudeService(limitRequest)")
    public Object doAround(ProceedingJoinPoint pjp, LimitRequest limitRequest) throws Throwable {
        // 获得request对象
        RequestAttributes ra = RequestContextHolder.getRequestAttributes();
        ServletRequestAttributes sra = (ServletRequestAttributes) ra;
        HttpServletRequest request = sra.getRequest();

        //校验IP次数不能大于总次数
        if(limitRequest.ipCount() > limitRequest.count()){
            return Result.fail(ResultCodeEnum.USER_REQUEST_API_ERROR);
        }
        if(limitRequest.isAppointIp() && limitRequest.ipCount() > 0){
            return appointIpCount(pjp,limitRequest,request);
        }
        // 获取Map对象， 如果没有则返回默认值
        // 第一个参数是key， 第二个参数是默认值
        ExpiringMap<String, Integer> map = book.getOrDefault(request.getRequestURI(), ExpiringMap.builder().variableExpiration().build());
        String mapKey = null;
        //是否指定IP
        if(limitRequest.isAppointIp()){
            mapKey = IpAddressUtil.getIpAdrress(request);
        }else {
            mapKey = request.getRequestURI();
        }
        String requestURI = request.getRequestURI();
        Integer uCount = map.getOrDefault(mapKey, 0);
        log.info("doAround 接口限制拦截IP: " + mapKey + ",url: " + requestURI);
        if (uCount >= limitRequest.count()) { // 超过次数，不执行目标方法
            //这里的返回对象类型根据controller方法的返回方式一致
            return Result.fail(ResultCodeEnum.RATE_LIMIT_WARN);
        } else if (uCount == 0){ // 第一次请求时，设置开始有效时间
            map.put(mapKey, uCount + 1, ExpirationPolicy.CREATED, limitRequest.time(), TimeUnit.MILLISECONDS);
        } else { // 未超过次数， 记录数据加一
            map.put(mapKey, uCount + 1);
        }
        book.put(requestURI, map);

        // result的值就是被拦截方法的返回值
        Object result = pjp.proceed();
        return result;
    }

    //指定IP请求次数限制
    public Object appointIpCount(ProceedingJoinPoint pjp, LimitRequest limitRequest,HttpServletRequest request)throws Throwable{

        ExpiringMap<String, Integer> map1 = bookMap1.getOrDefault(request.getRequestURI(), ExpiringMap.builder().variableExpiration().build());
        ExpiringMap<String, Integer> map2 = bookMap2.getOrDefault(request.getRequestURI(), ExpiringMap.builder().variableExpiration().build());

        String requestURI = request.getRequestURI();

        String mapKey = request.getRequestURI() + IpAddressUtil.getIpAdrress(request);

        //获取路径对应的请求次数
        Integer urlCount = map1.getOrDefault(requestURI, 0);
        //获取IP请求对应的次数
        Integer ipCount = map2.getOrDefault(mapKey, 0);

        log.info("appointIpCount 接口限制拦截指定IP次数: " + ipCount + ",url: " + urlCount);

        if (urlCount >= limitRequest.count() || ipCount >= limitRequest.ipCount()) {
            return Result.fail(ResultCodeEnum.RATE_LIMIT_WARN);
        }
        if(urlCount == 0){
            map1.put(requestURI, urlCount + 1, ExpirationPolicy.CREATED, limitRequest.time(), TimeUnit.MILLISECONDS);
        }else {
            map1.put(requestURI, urlCount + 1);
        }
        if(ipCount == 0){
            map2.put(mapKey, ipCount + 1, ExpirationPolicy.CREATED, limitRequest.time(), TimeUnit.MILLISECONDS);
        }else {
            map2.put(mapKey,ipCount + 1);
        }

        bookMap1.put(requestURI, map1);
        bookMap2.put(requestURI, map2);

        Object result = pjp.proceed();
        return result;
    }


}
